#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GPU服务器自动注册模块
负责向API网关注册和心跳监测
"""

import requests
import time
import threading
import socket
import logging
import os
from config import Config

logger = logging.getLogger(__name__)

class AutoRegister:
    def __init__(self, gateway_url: str = None):
        self.gateway_url = gateway_url or Config.GATEWAY_URL
        self.internal_ip = Config.INTERNAL_IP
        self.public_ip = Config.PUBLIC_IP or self._get_public_ip()
        self.region = Config.REGION
        self.server_url = f"http://{self.internal_ip}:{Config.SERVER_PORT}"
        self.public_url = f"http://{self.public_ip}:{Config.SERVER_PORT}" if self.public_ip else None
        self.registered = False
        self.heartbeat_thread = None
        self.registration_thread = None
        self.running = False
        
        logger.info(f"GPU服务器初始化: 内网IP={self.internal_ip}, 公网IP={self.public_ip}, 地域={self.region}")
    
    def _get_internal_ip(self) -> str:
        """获取内网IP地址"""
        try:
            # 创建UDP socket连接来获取本机IP
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            internal_ip = s.getsockname()[0]
            s.close()
            return internal_ip
        except Exception as e:
            logger.warning(f"获取内网IP失败: {e}, 尝试备用方法")
            try:
                hostname = socket.gethostname()
                return socket.gethostbyname(hostname)
            except Exception as e2:
                logger.error(f"获取内网IP失败: {e2}")
                return "127.0.0.1"
    
    def _get_public_ip(self) -> str:
        """获取公网IP地址"""
        # 优先使用环境变量
        public_ip = os.getenv('GPU_SERVER_PUBLIC_IP')
        if public_ip:
            return public_ip
        
        # 尝试通过外部服务获取公网IP
        try:
            response = requests.get('https://api.ipify.org', timeout=3)
            if response.status_code == 200:
                return response.text.strip()
        except Exception as e:
            logger.warning(f"无法获取公网IP: {e}")
        
        return None
    
    def _detect_region(self) -> str:
        """根据内网IP自动检测地域（已由Config处理，保留用于向后兼容）"""
        return Config.REGION
        
    def _get_server_url(self) -> str:
        """获取服务器URL（保持向后兼容）"""
        return self.server_url
    
    def _get_gpu_info(self) -> dict:
        """获取GPU信息"""
        try:
            import torch
            if torch.cuda.is_available():
                return {
                    'gpu_count': torch.cuda.device_count(),
                    'gpu_name': torch.cuda.get_device_name(),
                    'memory_total': torch.cuda.get_device_properties(0).total_memory,
                    'available': True
                }
        except ImportError:
            pass
        
        return {
            'gpu_count': 0,
            'gpu_name': 'CPU',
            'memory_total': 0,
            'available': False
        }
    
    def register_to_gateway(self) -> bool:
        """向API网关注册"""
        try:
            gpu_info = self._get_gpu_info()
            
            # 发送完整的网络信息：内网IP、公网IP、地域
            registration_data = {
                'server_url': self.server_url,  # 内网URL (优先)
                'public_url': self.public_url,  # 公网URL (备用)
                'internal_ip': self.internal_ip,
                'public_ip': self.public_ip,
                'region': self.region,
                'gpu_info': gpu_info,
                'server_type': 'gpu_server',
                'secret': Config.WEBHOOK_SECRET,
                'timestamp': time.time()
            }
            
            logger.info(f"向API网关注册: 内网={self.server_url}, 公网={self.public_url}, 地域={self.region}")
            
            response = requests.post(
                f"{self.gateway_url}/webhook/register",
                json=registration_data,
                timeout=10
            )
            
            if response.status_code == 200:
                result = response.json()
                server_id = result.get('server_id', 'unknown')
                logger.info(f"成功注册到API网关: {self.gateway_url}, Server ID: {server_id}")
                self.registered = True
                return True
            else:
                logger.error(f"注册失败: {response.status_code} - {response.text}")
                return False
                
        except requests.exceptions.ConnectionError:
            logger.error(f"无法连接到API网关: {self.gateway_url}")
            return False
        except Exception as e:
            logger.error(f"注册异常: {str(e)}")
            return False
    
    def unregister_from_gateway(self) -> bool:
        """从API网关注销"""
        try:
            unregistration_data = {
                'server_url': self.server_url,
                'timestamp': time.time()
            }
            
            response = requests.post(
                f"{self.gateway_url}/webhook/unregister",
                json=unregistration_data,
                timeout=10
            )
            
            if response.status_code == 200:
                logger.info("成功从API网关注销")
                self.registered = False
                return True
            else:
                logger.error(f"注销失败: {response.status_code} - {response.text}")
                return False
                
        except Exception as e:
            logger.error(f"注销异常: {str(e)}")
            return False
    
    def send_heartbeat(self) -> bool:
        """发送心跳"""
        try:
            heartbeat_data = {
                'server_url': self.server_url,
                'status': 'healthy',
                'timestamp': time.time(),
                'gpu_info': self._get_gpu_info()
            }
            
            response = requests.post(
                f"{self.gateway_url}/webhook/heartbeat",
                json=heartbeat_data,
                timeout=5
            )
            
            if response.status_code == 200:
                logger.debug("心跳发送成功")
                return True
            else:
                logger.warning(f"心跳发送失败: {response.status_code}")
                return False
                
        except Exception as e:
            logger.warning(f"心跳发送异常: {str(e)}")
            return False
    
    def _heartbeat_loop(self):
        """心跳循环"""
        logger.info("开始心跳监测...")
        
        while self.running:
            try:
                if self.registered:
                    self.send_heartbeat()
                time.sleep(30)  # 每30秒发送一次心跳
                
            except Exception as e:
                logger.error(f"心跳循环异常: {str(e)}")
                time.sleep(30)
    
    def _registration_loop(self):
        """注册循环"""
        logger.info("开始自动注册...")
        
        while self.running:
            try:
                if not self.registered:
                    if self.register_to_gateway():
                        logger.info("注册成功")
                    else:
                        logger.warning("注册失败，30秒后重试")
                        time.sleep(30)
                else:
                    time.sleep(60)  # 已注册，每分钟检查一次
                    
            except Exception as e:
                logger.error(f"注册循环异常: {str(e)}")
                time.sleep(30)
    
    def start_registration(self):
        """启动自动注册"""
        self.running = True
        
        # 启动注册线程
        self.registration_thread = threading.Thread(target=self._registration_loop)
        self.registration_thread.daemon = True
        self.registration_thread.start()
        
        logger.info("自动注册已启动")
    
    def start_heartbeat(self):
        """启动心跳监测"""
        self.running = True
        
        # 启动心跳线程
        self.heartbeat_thread = threading.Thread(target=self._heartbeat_loop)
        self.heartbeat_thread.daemon = True
        self.heartbeat_thread.start()
        
        logger.info("心跳监测已启动")
    
    def stop(self):
        """停止所有服务"""
        logger.info("停止自动注册和心跳监测...")
        
        self.running = False
        
        # 注销
        if self.registered:
            self.unregister_from_gateway()
        
        # 等待线程结束
        if self.registration_thread and self.registration_thread.is_alive():
            self.registration_thread.join(timeout=5)
        
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.heartbeat_thread.join(timeout=5)
        
        logger.info("自动注册和心跳监测已停止")
